Repository navigation
Unthunk the cotangent in reduce(vcat/hcat) and hvcat pullbacks - #861
Open
CarloLucibello wants to merge 1 commit into
Open
CarloLucibello wants to merge 1 commit into
CarloLucibello wants to merge 1 commit into
Conversation
These pullbacks index the cotangent once per input array. Indexing an `AbstractThunk` forwards to `unthunk`, which recomputes the whole array on every access, so the cost was O(length(As) * length(dY)). Unthunk once at the top of each pullback, as the other `cat` pullbacks already do. Also bumps the patch version. Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #861 +/- ##
=======================================
Coverage 96.63% 96.64%
=======================================
Files 28 28
Lines 3598 3602 +4
=======================================
+ Hits 3477 3481 +4
Misses 121 121 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
devmotion
reviewed
Nov 2, 2026
devmotion
left a comment
devmotion
left a comment
Member
There was a problem hiding this comment.
Thanks, the fix looks good. Unthunking once at the top matches hcat/vcat/cat/stack. It leaves plain-array and ZeroTangent cotangents alone, and it handles InplaceableThunk.
Two suggestions on the tests:
- Rather than using
@allocatedthresholds as a stand-in, consider counting unthunks directly. That checks the property itself, gives the same result on every Julia version and coverage setting, and doesn't need the 1600-input or 200-argumenthvcatcases:This would replace the "cost does not grow with the number of inputs" testset and the@testset "$name" for (name, (Y, back)) in cases dY = rand(size(Y)...) expected = back(dY) n = Ref(0) counted() = (n[] += 1; copy(dY)) for dȲ in (@thunk(counted()), InplaceableThunk(dx -> dx .+= dY, @thunk(counted()))) n[] = 0 @test back(dȲ) == expected @test n[] == 1 end end
allocs/thunk/inplaceablehelpers. plain(dY) = dYcan go;expected = back(dY)does the same.
Possible follow-up: hvcat_pullback could return one InplaceableThunk per argument, as hcat/vcat/cat do, so slices for non-differentiable arguments are never computed.
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Closes #860.
The pullbacks of
reduce(vcat, As),reduce(hcat, As)(reduce_vcat_pullback,reduce_hcat_pullback_1,reduce_hcat_pullback_2) andhvcat(hvcat_pullback) index the cotangent once per input. Indexing anAbstractThunkcallsunthunkeach time, so a thunked cotangent (e.g. fromsum's pullback in Zygote, FluxML/Zygote.jl#1686) is recomputedlength(As)times. This callsunthunkonce at the top of each, ashcat/vcat/cat/stackalready do.I used
dY = unthunk(ȳ)(renaming the argument) rather thandY = unthunk(dY), matching the neighbouring pullbacks and avoiding rebinding a variable captured by thedoblocks.hvcat_pullbackhas the same pattern and is included.Measured on Julia 1.13.2 (min of 5 runs), thunked cotangent
@thunk(ones(size(Y))):reduce(vcat), 1600 ×Vector(10)InplaceableThunkreduce(vcat), 800 ×Matrix(10×2)reduce(hcat), 1600 ×Vector(10)reduce(hcat), 800 ×Matrix(10×2)hvcat, 200 ×Vector(10)Added a testset in
test/rulesets/Base/array.jlchecking that@thunkandInplaceableThunkcotangents give the same result as the plain array for all the variants above (vectors / matrices / mixed), and that@allocated back(thunk) < 10 * sizeof(Y)at 1600 inputs (the quadratic path allocates aboutlength(As) * sizeof(Y)). Without the fix, the 10 allocation checks fail and the 14 equality checks pass; with it all 24 pass, also under--check-bounds=yes --code-coverage=user.Tests:
test/runtests.jl rulesets/Baseon Julia 1.13.2 passes (16978 pass, 52 broken as before). Not run locally on 1.10 (left to CI).Bumps the patch version to 1.74.10, as recent PRs do.
Noticed but left alone:
vectpullbacks indexȳ[n]per element and are equally quadratic in the number of arguments (N=256: 0.8 ms vs 3 µs), but N is a compile-time argument count, so I did not touch them.src/rulesets/Base/either alreadyunthunkor access the cotangent once.🧔♂️🤖 Assisted by AI